/**
 * @author luzeyu & liangqiaowei
 */
package tiei.ajp.huffman.utils;

import org.junit.Test;

import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;

public class HuffmanEncoder<T> {
    private Map<T, String> codeMap = null;
    private Map<String, T> decodeMap = null;  // Maybe useless

    public HuffmanEncoder() {
        this.codeMap = new HashMap<>();
        this.decodeMap = new HashMap<>();
    }

    private void generateCodes(TreeNode<T> node, StringBuilder prefix) {
        // --- Recursion end condition: leaf node
        assert node != null;
        if (node.left == null && node.right == null) {
            assert node.data != null;
            assert prefix != null;
            codeMap.put(node.data, prefix.toString());
            decodeMap.put(prefix.toString(), node.data);
            return;
        }

        if (node.left != null) {
            prefix.append('0');
            generateCodes(node.left, prefix);
            prefix.deleteCharAt(prefix.length() - 1);
        }
        if (node.right != null) {
            prefix.append('1');
            generateCodes(node.right, prefix);
            prefix.deleteCharAt(prefix.length() - 1);
        }
    }

    /**
     * Generate byte encodings according to huffman tree
     *
     * @param huffmanTreeRoot tree
     * @return encoding map
     */
    public Map<T, String> encode(TreeNode<T> huffmanTreeRoot) {
        assert huffmanTreeRoot != null;
        generateCodes(huffmanTreeRoot, new StringBuilder());
        return codeMap;
    }

    public Map<String, T> getDecodeMap() {
        for (Map.Entry<T, String> entry : codeMap.entrySet()) {
            decodeMap.put(entry.getValue(), entry.getKey());
        }
        return decodeMap;
    }


//    private void linearize(TreeNode<T> node, StringBuilder sb) {
//        assert node != null;
//
//        if (node.left == null && node.right == null) sb.append('0').append(node.data);
//        else {
//            sb.append('1');
//            linearize(node.left, sb);
//            linearize(node.right, sb);
//        }
//    }
//
//    public String linearize(TreeNode<T> root) {
//        StringBuilder sb = new StringBuilder();
//        linearize(root, sb);
//        return sb.toString();
//    }
//
//    public TreeNode<T> delinearize(String s) {
//        return delinearize(s, new int[]{0});
//    }
//
//    private TreeNode<T> delinearize(String s, int[] index) {
//        if (index[0] >= s.length()) return null;
//        char c = s.charAt(index[0]++);
//        if (c == '0') {
//            T data = (T) Character.valueOf(s.charAt(index[0]++));
//            return new TreeNode<T>(data);
//        } else if (c == '1') {
//            TreeNode<T> node = new TreeNode<>();
//            node.left = delinearize(s, index);
//            node.right = delinearize(s, index);
//            return node;
//        }
//        return null;
//    }

    @Test
    public void testEncode() {
        System.out.println("test");
        TreeNode<Character> root = new TreeNode<>(15, null);
        TreeNode<Character> rl = new TreeNode<>(7, null);
        TreeNode<Character> rr = new TreeNode<>(8, 'd');
        TreeNode<Character> rll = new TreeNode<>(3, null);
        TreeNode<Character> rlr = new TreeNode<>(4, 'c');
        TreeNode<Character> rlll = new TreeNode<>(1, 'a');
        TreeNode<Character> rllr = new TreeNode<>(2, 'b');

        root.left = rl;
        root.right = rr;
        rl.left = rll;
        rl.right = rlr;
        rll.left = rlll;
        rll.right = rllr;

        HuffmanEncoder<Character> he = new HuffmanEncoder<>();
        Map<Character, String> encMap = he.encode(root);
        assert encMap.get('d').equals("1");
        assert encMap.get('c').equals("01");
        assert encMap.get('b').equals("001");
        assert encMap.get('a').equals("000");

        System.out.println("Test pass");
    }

}

